/*
 * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
 * with the License. A copy of the License is located at
 *
 * http://aws.amazon.com/apache2.0/
 *
 * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
 * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
 * and limitations under the License.
 */
package com.yhzdys.ocr.service.support;

import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.index.NDIndex;
import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;
import ai.djl.translate.NoBatchifyTranslator;
import ai.djl.translate.TranslatorContext;
import ai.djl.util.Utils;

import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.List;

/**
 * A {@link PpWordRecognitionTranslator} that preprocess {@link Image} post-process the {@link NDArray} into text.
 */
public class PpWordRecognitionTranslator implements NoBatchifyTranslator<Image, String> {

    private volatile List<String> table;

    /**
     * {@inheritDoc}
     */
    @Override
    public void prepare(TranslatorContext ctx) throws IOException {
        if (table != null) {
            return;
        }
        synchronized (this) {
            if (table != null) {
                return;
            }
            try (InputStream inputStream = ctx.getModel().getArtifact("ppocr_keys_v1.txt").openStream()) {
                List<String> list = Utils.readLines(inputStream, true);
                list.add(0, "blank");
                list.add("");
                table = list;
            }
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public NDList processInput(TranslatorContext ctx, Image input) {
        NDManager ctxNDManager = ctx.getNDManager();
        NDArray image = input.toNDArray(ctxNDManager, Image.Flag.COLOR);

        int imgC = 3, imgH = 48, imgW = 320;
        float max_wh_ratio = (float) imgW / (float) imgH;
        int h = input.getHeight();
        int w = input.getWidth();
        float wh_ratio = (float) w / (float) h;
        max_wh_ratio = Math.max(max_wh_ratio, wh_ratio);
        imgW = (int) (imgH * max_wh_ratio);
        int resized_w;
        if (Math.ceil(imgH * wh_ratio) > imgW) {
            resized_w = imgW;
        } else {
            resized_w = (int) (Math.ceil(imgH * wh_ratio));
        }

        NDArray resizedImage = NDImageUtils.resize(image, resized_w, imgH);
        NDArray transposedImage = resizedImage.transpose(2, 0, 1).toType(DataType.FLOAT32, false);
        transposedImage.divi(255F).subi(0.5F).divi(0.5F);
        NDArray paddingIM = ctxNDManager.zeros(new Shape(imgC, imgH, imgW), DataType.FLOAT32);
        paddingIM.set(new NDIndex(":,:,0:" + resized_w), transposedImage);
        NDArray paddingIM0 = paddingIM.flip(0);
        NDArray paddingIM1 = paddingIM0.expandDims(0);
        return new NDList(paddingIM1);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public String processOutput(TranslatorContext ctx, NDList list) {
        List<String> table = this.table;
        StringBuilder builder = new StringBuilder();

        NDArray tokens = list.singletonOrThrow();
        NDArray ndArray = tokens.get(0);
        NDArray ndArray1 = ndArray.argMax(1);

        long[] indices = ndArray1.toLongArray();
        boolean[] selection = new boolean[indices.length];
        Arrays.fill(selection, true);
        for (int i = 1; i < indices.length; i++) {
            if (indices[i] == indices[i - 1]) {
                selection[i] = false;
            }
        }
        int lastIdx = 0;
        for (int i = 0; i < indices.length; i++) {
            if (selection[i] && indices[i] > 0 && !(i > 0 && indices[i] == lastIdx)) {
                builder.append(table.get((int) indices[i]));
            }
        }
        return builder.toString();
    }
}
